from rest_framework.renderers import JSONRenderer


class CustomRenderer(JSONRenderer):
    def render(self, data, accepted_media_type=None, renderer_context=None):
        # 获取响应状态码
        status_code = renderer_context['response'].status_code

        # 如果响应数据中存在 'code'，'message'，'data' 字段，则不做处理
        if data and 'code' in data and 'message' in data and 'data' in data:
            return super().render(data, accepted_media_type, renderer_context)

        # 根据响应状态码设置 'code' 和 'message' 字段
        if status_code >= 400:
            code = 1
            message = '请求失败'
        else:
            code = 0
            message = '请求成功'

        # 包装响应数据
        wrapped_data = {
            'code': code,
            'message': message,
            'data': data
        }

        return super().render(wrapped_data, accepted_media_type, renderer_context)
